<<<<<<< HEAD
from BC_learning.Agent_utils import BC_Agent
=======
from BC_learning.Agent_utils import EGPO_Agent, CatObs
>>>>>>> 181f11012f613eb466d33564fc88ba83963d4e05
from hex_cfg import HexCfg
import torch
import os
import rospy
from std_msgs.msg import Float64MultiArray
from interface.msg import joy_command

class motor_mode:
    Disable = -1
    Zero_Torque = 0
    Pos_vel = 1
    Traj_follow = 2


def ProcessCommand(msg:joy_command):
    global agent
    if msg.disable_torque:
        for i in range(6):
            q_des_msg_list[i].data=[motor_mode.Disable,0,0,0,0,0,0,-999]
            q_des_pub_list[i].publish(q_des_msg_list[i])
        return

<<<<<<< HEAD
    command=torch.tensor([[msg.set_init,msg.x_vec,msg.y_vec,0,msg.w_twist]],dtype=torch.float32,device=device)
    command[:,1]=command[:,1]*hex_cfg.max_vec.x
    command[:,2]=command[:,2]*hex_cfg.max_vec.y
    command[:,3]=command[:,3]*hex_cfg.max_vec.z
    command[:,4]=command[:,4]*hex_cfg.max_vec.omega_move    
    obs=torch.cat([q_cur.reshape(1,18),q_dot_cur.reshape(1,18),q_torq.reshape(1,18),command[:,1:5]],dim=1)
=======
    command=torch.tensor([[msg.set_init,msg.x_vec,msg.y_vec,msg.w_twist]],dtype=torch.float32,device=device)
    command[:,1]=command[:,1]* 0.5 * 2.0
    command[:,2]=command[:,2]* 0.8 * 2.0
    command[:,3]=command[:,3]* 2.0 * 0.25
    #拼接观测，要乘上观测值对应的系数
    obs = torch.cat([
        quat,
        omega * 0.25,
        acc,
        (q_cur-q_default).reshape(1,-1),
        q_dot_cur.reshape(1,-1) * 0.05,
        q_torq.reshape(1,-1) * 0.1,
        command[:,1:]
    ],dim=1)
>>>>>>> 181f11012f613eb466d33564fc88ba83963d4e05
    # print("obs.shape=",obs.shape)

    with torch.no_grad():
        action=agent(obs).reshape(6,3)
        #还原到关节期望角度
        action = action*0.5+q_default
    for i in range(6):
        if msg.set_init==0:
            q_des_msg_list[i].data=[motor_mode.Traj_follow,action[i,0].item(),action[i,1].item(),action[i,2].item(),0,0,0,-999]
        else:
            q_des_msg_list[i].data=[motor_mode.Pos_vel,q_default[i,0],q_default[i,1],q_default[i,2],0,0,0,-999]
        q_des_pub_list[i].publish(q_des_msg_list[i])


def UpdateQState(q_state_msgs:Float64MultiArray,index):
    global device, q_cur, q_dot_cur, q_torq
    q_cur[:,index,0:3]=torch.tensor(q_state_msgs.data[0:3],dtype=torch.float32,device=device)
    q_torq[:,index,0:3]=torch.tensor(q_state_msgs.data[3:6],dtype=torch.float32,device=device)
    q_dot_cur[:,index,0:3]=torch.tensor(q_state_msgs.data[6:9],dtype=torch.float32,device=device)

def UpdateRootState(root_state_msgs:Float64MultiArray):
    global device, acc, omega, quat
    quat[0]=torch.tensor(root_state_msgs.data[0:4],dtype=torch.float32,device=device)
    omega[0]=torch.tensor(root_state_msgs.data[4:7],dtype=torch.float32,device=device)
    acc[0]=torch.tensor(root_state_msgs.data[7:10],dtype=torch.float32,device=device)

    print("quat=",quat)
    print("omega=",omega)
    print("acc=",acc)

if __name__ == '__main__':
    rospy.init_node('run_agent', anonymous=True)
    hex_cfg=HexCfg("---")
    command_sub=rospy.Subscriber('/usr/command',joy_command,ProcessCommand,queue_size=10)
    imu_sub = rospy.Subscriber('/imu/model_states',Float64MultiArray,UpdateRootState,queue_size=10)


    device='cpu'
    q_default = torch.zeros(6,3,dtype=torch.float32,device=device)
    q_des=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    q_cur=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    q_dot_cur=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    q_torq=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    #填充关节默认角度
    for i in range(6):
        if i==0 or i==4:
            thigh=0.35
        elif i==1 or i==3:
            thigh=-0.35
        else:
            thigh=0.0
        knee=0.7
        ankle=-2.14
        q_default[i,0]=thigh
        q_default[i,1]=knee
        q_default[i,2]=ankle
        
    #IMU观测
    acc = torch.zeros(1,3,dtype=torch.float32,device=device)
    omega = torch.zeros(1,3,dtype=torch.float32,device=device)
    quat = torch.zeros(1,4,dtype=torch.float32,device=device)
    
<<<<<<< HEAD
    agent=BC_Agent()
    agent.load_state_dict(torch.load(os.getcwd()+'/bag/dagger_agent_omni.pth'))
=======
    # agent=BC_Agent()
    agent = EGPO_Agent()
    # agent.load_state_dict(torch.load(os.getcwd()+'/bag/dagger_agent.pth',weights_only=True))
    agent.load_state_dict(torch.load(os.getcwd()+'/src/agent/EGPO_actor.pt',weights_only=True))
>>>>>>> 181f11012f613eb466d33564fc88ba83963d4e05
    agent.eval()
    agent.to(device)


    q_des_pub_list=[]
    q_cur_sub_list=[]
    q_des_msg_list=[] #Float64MultiArray()
    q_state_dict={'q_des':[],'q_cur':[],'q_dot_cur':[],'torq_cur':[]}
    leg_name_list=['LB','LF','LM','RB','RF','RM']
    for i,leg_name in enumerate(leg_name_list):
        q_des_pub_list.append(rospy.Publisher('/'+leg_name+'/sita_des',Float64MultiArray,queue_size=10))
        q_cur_sub_list.append(rospy.Subscriber('/'+leg_name+'/sita_cur',Float64MultiArray,UpdateQState,callback_args=i))
        q_des_msg_list.append(Float64MultiArray())
    print("------------------->run curv adapt ready<-------------------")
    print(f" device ={device} ; load agent from {os.getcwd()+'/src/agent/EGPO_actor.pt'}")
    rospy.spin()